import requests
import json
import time

# 测试配置
TEST_PORTS = [8001]  # 要测试的端口列表 (根据终端信息，服务在8001端口运行)
TEST_ENDPOINTS = [
    '/',
    '/api-info',
    '/test/document-count',
    '/documents'
]  # 要测试的端点

# 认证信息（如果需要）
AUTH_CREDENTIALS = {
    'username': 'admin',
    'password': 'admin'  # 替换为实际的管理员密码
}

# 是否需要认证
REQUIRES_AUTH = False

# 测试文档内容
TEST_DOCUMENT = """
这是一个测试文档，用于诊断问答系统问题。
- 客服电话: 400-123-4567
- 服务时间: 周一至周五 9:00-18:00
- 常见问题: 如何修改密码、如何退款、如何投诉等
- 投诉邮箱: service@example.com
"""

TEST_DOCUMENT_NAME = "test_document_management.txt"

# 获取认证Token
def get_auth_token(port):
    url = f'http://localhost:{port}/token'
    try:
        response = requests.post(url, data=AUTH_CREDENTIALS)
        if response.status_code == 200:
            return response.json().get('access_token')
        else:
            print(f'获取Token失败: {response.status_code} {response.reason}')
            return None
    except Exception as e:
        print(f'获取Token异常: {str(e)}')
        return None

# 上传测试文档
def upload_test_document(port, token):
    url = f'http://localhost:{port}/upload-document/'
    headers = {'Authorization': f'Bearer {token}'} if token else {}
    
    # 创建测试文件
    files = {'file': (TEST_DOCUMENT_NAME, TEST_DOCUMENT)}
    
    try:
        response = requests.post(url, headers=headers, files=files)
        print(f'文档上传状态: {response.status_code} {response.reason}')
        
        if response.status_code == 200:
            try:
                data = response.json()
                print(f'  上传成功: {data.get("filename")}, 块数: {data.get("vectorization", {}).get("chunk_count")}')
                return True
            except ValueError:
                print('  无法解析上传响应')
        return False
    except Exception as e:
        print(f'  上传错误: {str(e)}')
        return False

# 删除测试文档
def delete_test_document(port, token):
    url = f'http://localhost:{port}/documents/{TEST_DOCUMENT_NAME}'
    headers = {'Authorization': f'Bearer {token}'} if token else {}
    
    try:
        response = requests.delete(url, headers=headers)
        print(f'文档删除状态: {response.status_code} {response.reason}')
        
        if response.status_code == 200:
            try:
                data = response.json()
                print(f'  删除结果: {data.get("message")}')
                return data.get('success', False)
            except ValueError:
                print('  无法解析删除响应')
        return False
    except Exception as e:
        print(f'  删除错误: {str(e)}')
        return False

# 测试RAG查询
def test_rag_query(port, token):
    url = f'http://localhost:{port}/rag-query'
    headers = {'Authorization': f'Bearer {token}'} if token else {}
    headers['Content-Type'] = 'application/json'
    
    # 测试问题列表
    queries = [
        "客服电话是多少?",
        "服务时间是什么时候?",
        "投诉邮箱是什么?",
        "如何修改密码?"
    ]
    
    session = requests.Session()
    
    for query in queries:
        print(f'\n测试查询: "{query}"')
        try:
            response = session.post(
                url,
                headers=headers,
                json={'query': query}
            )
            print(f'  查询状态: {response.status_code} {response.reason}')
            
            if response.status_code == 200:
                try:
                    data = response.json()
                    print(f'  回答: {data.get("answer")[:100]}...')
                    print(f'  引用文档数: {len(data.get("references", []))}')
                    
                    # 检查是否有引用文档
                    if len(data.get("references", [])) == 0:
                        print('  警告: 未找到相关文档引用')
                    else:
                        # 打印第一个引用的来源
                        ref = data['references'][0]
                        print(f'  引用来源: {ref.get("source")} (相似度: {ref.get("similarity", 0):.2f})')
                except ValueError:
                    print('  无法解析查询响应')
        except Exception as e:
            print(f'  查询错误: {str(e)}')

# 运行测试
def run_tests():
    print('开始详细端口测试...')
    print('=' * 60)

    for port in TEST_PORTS:
        print(f'\n=== 测试端口: {port} ===')
        token = get_auth_token(port)
        headers = {'Authorization': f'Bearer {token}'} if token else {}
        session = requests.Session()

        # 测试基本端点
        print('\n--- 基本端点测试 ---')
        for endpoint in TEST_ENDPOINTS:
            url = f'http://localhost:{port}{endpoint}'
            print(f'测试端点: {url}')

            try:
                response = session.get(url, headers=headers, timeout=5)
                print(f'状态: {response.status_code} {response.reason}')

                # 尝试解析JSON响应
                try:
                    data = response.json()
                    print(f'  响应数据长度: {len(json.dumps(data))} 字节')
                    # 对于文档计数端点，打印详细信息
                    if endpoint == '/test/document-count':
                        print(f'  文档数量: {data.get("document_count")}')
                        print(f'  原始文本块数: {data.get("raw_text_chunks_count")}')
                        print(f'  原始元数据数: {data.get("raw_metadata_count")}')
                    # 对于文档列表，打印数量和详情
                    elif endpoint == '/documents' and 'documents' in data:
                        print(f'  文档数量: {len(data["documents"])}')
                        for doc in data["documents"]:
                            print(f'  - {doc.get("source")} ({doc.get("chunk_count")}个块)')
                except ValueError:
                    print(f'  无法解析JSON响应，内容长度: {len(response.text)} 字节')
            except Exception as e:
                print(f'  错误: {str(e)}')

        # 文档管理测试
        print('\n--- 文档管理测试 ---')
        # 先检查是否有测试文档，如果有则删除
        doc_list_url = f'http://localhost:{port}/documents'
        try:
            # 根据是否需要认证决定是否添加headers
            request_headers = headers if REQUIRES_AUTH else {}
            response = session.get(doc_list_url, headers=request_headers)
            if response.status_code == 200:
                docs = response.json().get('documents', [])
                test_doc_exists = any(doc.get('source') == TEST_DOCUMENT_NAME for doc in docs)
                if test_doc_exists:
                    print(f'  测试文档 {TEST_DOCUMENT_NAME} 已存在，尝试删除...')
                    delete_test_document(port, token if REQUIRES_AUTH else None)
        except Exception as e:
            print(f'  检查文档列表错误: {str(e)}')

        # 上传测试文档
        print('  上传测试文档...')
        upload_success = upload_test_document(port, token if REQUIRES_AUTH else None)

        # 再次获取文档列表
        if upload_success:
            print('  上传后获取文档列表...')
            try:
                response = session.get(doc_list_url, headers=request_headers)
                if response.status_code == 200:
                    docs = response.json().get('documents', [])
                    print(f'  上传后文档数量: {len(docs)}')
                    test_doc_exists = any(doc.get('source') == TEST_DOCUMENT_NAME for doc in docs)
                    print(f'  测试文档存在: {test_doc_exists}')
            except Exception as e:
                print(f'  获取文档列表错误: {str(e)}')

        # 测试RAG查询
        print('\n--- RAG查询测试 ---')
        if upload_success:
            test_rag_query(port, token if REQUIRES_AUTH else None)
        else:
            print('  文档上传失败，跳过RAG查询测试')

        # 测试删除文档
        print('\n--- 删除文档测试 ---')
        if upload_success:
            print('  删除测试文档...')
            delete_success = delete_test_document(port, token if REQUIRES_AUTH else None)

            # 验证删除
            if delete_success:
                print('  验证删除结果...')
                try:
                    response = session.get(doc_list_url, headers=request_headers)
                    if response.status_code == 200:
                        docs = response.json().get('documents', [])
                        print(f'  删除后文档数量: {len(docs)}')
                        test_doc_exists = any(doc.get('source') == TEST_DOCUMENT_NAME for doc in docs)
                        print(f'  测试文档仍存在: {test_doc_exists}')
                except Exception as e:
                    print(f'  验证删除错误: {str(e)}')

        # 测试会话管理
        print('\n--- 会话管理测试 ---')
        if token:
            # 测试RAG查询
            rag_url = f'http://localhost:{port}/rag-query'
            print(f'测试RAG查询: {rag_url}')
            try:
                response = session.post(
                    rag_url,
                    headers={**headers, 'Content-Type': 'application/json'},
                    json={'query': '测试问题'}
                )
                print(f'  状态: {response.status_code} {response.reason}')
                # 检查Cookie
                cookies = session.cookies.get_dict()
                print(f'  会话Cookie: {cookies}')
            except Exception as e:
                print(f'  错误: {str(e)}')

        print('\n--- 跨端口Cookie测试 ---')
        # 测试从一个端口获取的Cookie在另一个端口是否可用
        if len(TEST_PORTS) > 1:
            other_port = TEST_PORTS[1] if port == TEST_PORTS[0] else TEST_PORTS[0]
            print(f'使用端口 {port} 的Cookie访问端口 {other_port}')
            try:
                # 使用当前会话的Cookie访问其他端口
                cross_url = f'http://localhost:{other_port}/api-info'
                response = session.get(cross_url, timeout=5)
                print(f'  状态: {response.status_code} {response.reason}')
            except Exception as e:
                print(f'  错误: {str(e)}')
        else:
            print('  只有一个测试端口，跳过跨端口测试')

    print('
\n' + '=' * 60)
    print('测试完成')

if __name__ == '__main__':
    run_tests()